{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "from neo4j import GraphDatabase, RoutingControl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python -m spacy download zh_core_web_sm -i https://pypi.tuna.tsinghua.edu.cn/simple\n",
    "# !pip install neo4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "提取的人物： ['钮祜禄·甄嬛', '王爱慕', '王死讯', '滴血验亲']\n"
     ]
    }
   ],
   "source": [
    "# 加载Spacy模型\n",
    "nlp = spacy.load(\"zh_core_web_sm\")\n",
    "\n",
    "# 示例文本\n",
    "# text = \"甄嬛，原名甄玉嬛，嫌玉字俗气而改名甄嬛，为汉人甄远道之女，后被雍正赐姓钮祜禄氏，抬旗为满洲上三旗，获名“钮祜禄·甄嬛”。\"\n",
    "text = \"\"\"\n",
    "甄嬛，原名甄玉嬛，嫌玉字俗气而改名甄嬛，为汉人甄远道之女，后被雍正赐姓钮祜禄氏，抬旗为满洲上三旗，获名“钮祜禄·甄嬛”。\n",
    "同沈眉庄、安陵容参加选秀，因容貌酷似纯元皇后而被选中。入宫后面对华妃的步步紧逼，沈眉庄被冤、安陵容变心，从偏安一隅的青涩少女变成了能引起血雨腥风的宫斗老手。\n",
    "雍正发现年氏一族的野心后令其父甄远道剪除，甄嬛也于后宫中用她的连环巧计帮皇帝解决政敌，故而深得雍正爱待。几经周折，终于斗垮了嚣张跋扈的华妃。\n",
    "甄嬛封妃时遭皇后宜修暗算，被皇上嫌弃，生下女儿胧月后心灰意冷，自请出宫为尼。然得果郡王爱慕，二人相爱，得知果郡王死讯后立刻设计与雍正再遇，风光回宫。\n",
    "此后甄父冤案平反、甄氏复起，她也生下双生子，在滴血验亲等各种阴谋中躲过宜修的暗害，最后以牺牲自己亲生胎儿的方式扳倒了幕后黑手的皇后。\n",
    "但雍正又逼甄嬛毒杀允礼，以测试甄嬛真心，并让已经生产过孩子的甄嬛去准格尔和亲。甄嬛遂视皇帝为最该毁灭的对象，大结局道尽“人类的一切争斗，皆因统治者的不公不义而起”，并毒杀雍正。\n",
    "四阿哥弘历登基为乾隆，甄嬛被尊为圣母皇太后，权倾朝野，在如懿传中安度晚年。\n",
    "\"\"\"\n",
    "\n",
    "# 对文本进行NER\n",
    "doc = nlp(text)\n",
    "persons = [ent.text for ent in doc.ents if ent.label_ == \"PERSON\"]\n",
    "\n",
    "# 提取人物属性\n",
    "attributes = {\"甄嬛\": {\"属性\": [\"美丽\", \"聪明\"], \"父亲\": \"甄远孙\"}}\n",
    "\n",
    "# 关系抽取\n",
    "relations = {\"甄嬛-皇帝\": \"特殊关系\"}\n",
    "\n",
    "print(\"提取的人物：\", persons)\n",
    "# print(\"人物属性：\", attributes)\n",
    "# print(\"人物关系：\", relations)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neo4j\n",
    "from to 的关系是：from 的 relation 是 to 。from是主人公，to 是他的朋友、父亲等等。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Neo4j 数据库已被重置清空。\n"
     ]
    }
   ],
   "source": [
    "from neo4j import GraphDatabase, RoutingControl\n",
    "\n",
    "\n",
    "URI      = \"neo4j+s://5bb87071.databases.neo4j.io\"\n",
    "username = \"neo4j\"\n",
    "password = \"PPcEmlN5ognHV-YG0J9Puj0FzCwdiXsElSbaHugN9xA\"\n",
    "AUTH     = (username, password)\n",
    "\n",
    "def clear_all(driver):\n",
    "    # 删除所有节点和关系\n",
    "    driver.execute_query(\n",
    "        \"MATCH (n)\"\n",
    "        \"DETACH DELETE n\"\n",
    "    )\n",
    "    print('Neo4j 数据库已被重置清空。')\n",
    "\n",
    "def add_friend(driver, name, friend_name):\n",
    "    driver.execute_query(\n",
    "        \"MERGE (a:Person {name: $name}) \"\n",
    "        \"MERGE (friend:Person {name: $friend_name}) \"\n",
    "        \"MERGE (a)-[:KNOWS]->(friend)\",\n",
    "        name=name, friend_name=friend_name, database_=\"neo4j\",\n",
    "    )\n",
    "\n",
    "\n",
    "def print_friends(driver, name):\n",
    "    records, _, _ = driver.execute_query(\n",
    "        \"MATCH (a:Person)-[:KNOWS]->(friend) WHERE a.name = $name \"\n",
    "        \"RETURN friend.name ORDER BY friend.name\",\n",
    "        name=name, database_=\"neo4j\", routing_=RoutingControl.READ,\n",
    "    )\n",
    "    for record in records:\n",
    "        print(record[\"friend.name\"])\n",
    "\n",
    "\n",
    "with GraphDatabase.driver(URI, auth=AUTH) as driver:\n",
    "    # add_friend(driver, \"Arthur\", \"Guinevere\")\n",
    "    # add_friend(driver, \"Arthur\", \"Lancelot\")\n",
    "    # add_friend(driver, \"Arthur\", \"Merlin\")\n",
    "    clear_all(driver)\n",
    "    \n",
    "    print_friends(driver, \"Arthur\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Neo4j 数据库已被重置清空。\n",
      "name: 甄嬛, gender: 女, alias: 熹贵妃\n",
      "name: 雍正帝, gender: 男, alias: 皇上\n",
      "name: 皇后, gender: 女, alias: 景仁宫娘娘\n",
      "name: 华妃, gender: 女, alias: 华贵妃\n",
      "name: 沈眉庄, gender: 女, alias: 沈贵人\n",
      "name: 安陵容, gender: 女, alias: 安贵人\n",
      "name: 果郡王, gender: 男, alias: 爱新觉罗·允礼\n",
      "name: 温实初, gender: 男, alias: 太医\n",
      "name: 甄远道, gender: 男, alias: \n",
      "name: 胧月公主, gender: 女, alias: \n",
      "甄嬛的女儿是: 胧月公主\n",
      "甄嬛的别名: 熹贵妃\n",
      "人物信息:\n",
      "姓名: 甄嬛\n",
      "性别: 女\n",
      "别名: 熹贵妃\n"
     ]
    }
   ],
   "source": [
    "# 定义Neo4j会话\n",
    "driver = GraphDatabase.driver(URI, auth=(username, password))\n",
    "session = driver.session()\n",
    "\n",
    "clear_all(driver)\n",
    "\n",
    "# 导入人物属性\n",
    "with open('./Persons.csv', 'r', encoding='gbk') as file:\n",
    "    lines = file.readlines()[1:]  # 跳过CSV的表头行\n",
    "    for line in lines:\n",
    "        values = line.strip().split(',')\n",
    "        name, gender, alias = values\n",
    "        print(f'name: {name}, gender: {gender}, alias: {alias}')\n",
    "        cypher_query = f\"CREATE (:Character {{name: '{name}', gender: '{gender}', alias: '{alias}'}})\"\n",
    "        session.run(cypher_query)\n",
    "\n",
    "# 导入人物关系\n",
    "with open('./Relationship.csv', 'r', encoding='gbk') as file:\n",
    "    lines = file.readlines()[1:]  # 跳过CSV的表头行\n",
    "    for line in lines:\n",
    "        from_id, to_id, relation = line.strip().split(',')\n",
    "        cypher_query = f\"MATCH (from:Character {{name: '{from_id}'}}) \" \\\n",
    "                       f\"MATCH (to:Character {{name: '{to_id}'}}) \" \\\n",
    "                       f\"CREATE (from)-[:{relation}]->(to)\"\n",
    "        session.run(cypher_query)\n",
    "\n",
    "def query_alias(driver, person_name):\n",
    "    with driver.session() as session:\n",
    "        result = session.run(\"MATCH (person:Character {name: $person_name})\"\n",
    "                                \"RETURN person.alias AS alias\",\n",
    "                                person_name=person_name)\n",
    "        alias = result.single()[\"alias\"]\n",
    "\n",
    "        return alias\n",
    "        \n",
    "def query_relationship(driver, person_name, relationship_type):\n",
    "    with driver.session() as session:\n",
    "        query = f\"MATCH (person:Character {{name: '{person_name}'}})-[:{relationship_type}]->(related) RETURN related.name AS related_name\"\n",
    "        result = session.run(query)\n",
    "        related_names = [record[\"related_name\"] for record in result]\n",
    "        return related_names\n",
    "    \n",
    "def query_person_info(driver, person_name):\n",
    "    with driver.session() as session:\n",
    "        result = session.run(\"MATCH (person:Character {name: $person_name})\"\n",
    "                                \"RETURN person.name AS name, person.gender AS gender, person.alias AS alias\",\n",
    "                                person_name=person_name)\n",
    "        record = result.single()\n",
    "        if record:\n",
    "            return {\n",
    "                \"name\": record[\"name\"],\n",
    "                \"gender\": record[\"gender\"],\n",
    "                \"alias\": record[\"alias\"]\n",
    "            }\n",
    "        else:\n",
    "            return None\n",
    "    \n",
    "person_name = '甄嬛'\n",
    "relationship_type = '女儿'\n",
    "related_names = query_relationship(driver, person_name, relationship_type)\n",
    "if related_names:\n",
    "    print(f\"{person_name}的{relationship_type}是: {', '.join(related_names)}\")\n",
    "else:\n",
    "    print(f\"{person_name}没有{relationship_type}\")\n",
    "\n",
    "aliases = query_alias(driver, person_name)\n",
    "if aliases:\n",
    "    print(f\"{person_name}的别名: {aliases}\")\n",
    "else:\n",
    "    print(f\"{person_name}没有别名\")\n",
    "\n",
    "person_info = query_person_info(driver, person_name)\n",
    "if person_info:\n",
    "    print(\"人物信息:\")\n",
    "    print(f\"姓名: {person_info['name']}\")\n",
    "    print(f\"性别: {person_info['gender']}\")\n",
    "    print(f\"别名: {person_info['alias']}\")\n",
    "else:\n",
    "    print(f\"{person_name}不存在\")\n",
    "\n",
    "# 关闭会话和驱动\n",
    "session.close()\n",
    "driver.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 封装"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of nodes in the database: 10\n",
      "甄嬛的父亲是: 甄远道\n",
      "甄嬛的别名: 熹贵妃\n",
      "人物信息:\n",
      "姓名: 甄嬛\n",
      "性别: 女\n",
      "别名: 熹贵妃\n"
     ]
    }
   ],
   "source": [
    "from neo4j import GraphDatabase\n",
    "\n",
    "class Neo4jManager:\n",
    "    def __init__(self, uri, username, password):\n",
    "        self.driver = GraphDatabase.driver(uri, auth=(username, password))\n",
    "    \n",
    "    def close(self):\n",
    "        self.driver.close()\n",
    "\n",
    "    def import_characters(self, csv_path):\n",
    "        with open(csv_path, 'r', encoding='gbk') as file:\n",
    "            lines = file.readlines()[1:]\n",
    "            with self.driver.session() as session:\n",
    "                for line in lines:\n",
    "                    name, gender, alias = line.strip().split(',')\n",
    "                    session.run(\"CREATE (:Character {name: $name, gender: $gender, alias: $alias})\",\n",
    "                                name=name, gender=gender, alias=alias)\n",
    "\n",
    "    def import_relationships(self, csv_path):\n",
    "        with open(csv_path, 'r', encoding='gbk') as file:\n",
    "            lines = file.readlines()[1:]\n",
    "            with self.driver.session() as session:\n",
    "                for line in lines:\n",
    "                    from_id, to_id, relation = line.strip().split(',')\n",
    "                    cypher_query = f\"MATCH (from:Character {{name: '{from_id}'}}) \" \\\n",
    "                                f\"MATCH (to:Character {{name: '{to_id}'}}) \" \\\n",
    "                                f\"CREATE (from)-[:{relation}]->(to)\"\n",
    "                    session.run(cypher_query)\n",
    "\n",
    "    def clear_database(self):\n",
    "        with self.driver.session() as session:\n",
    "            session.run(\"MATCH (n) DETACH DELETE n\")\n",
    "\n",
    "    def query_status(self):\n",
    "        with self.driver.session() as session:\n",
    "            result     = session.run(\"MATCH (n) RETURN COUNT(n) AS node_count\")\n",
    "            node_count = result.single()[\"node_count\"]\n",
    "            return f\"Number of nodes in the database: {node_count}\"\n",
    "        \n",
    "    def query_alias(self, person_name):\n",
    "        with self.driver.session() as session:\n",
    "            result = session.run(\"MATCH (person:Character {name: $person_name})\"\n",
    "                                    \"RETURN person.alias AS alias\",\n",
    "                                    person_name=person_name)\n",
    "            alias = result.single()[\"alias\"]\n",
    "\n",
    "            return alias\n",
    "            \n",
    "    def query_relationship(self, person_name, relationship_type):\n",
    "        with self.driver.session() as session:\n",
    "            query = f\"MATCH (person:Character {{name: '{person_name}'}})-[:{relationship_type}]->(related) \\\n",
    "                        RETURN related.name AS related_name\"\n",
    "            result        = session.run(query)\n",
    "            related_names = [record[\"related_name\"] for record in result]\n",
    "            return related_names\n",
    "        \n",
    "    def query_person_info(self, person_name):\n",
    "        with self.driver.session() as session:\n",
    "            result = session.run(\"MATCH (person:Character {name: $person_name})\"\n",
    "                                    \"RETURN person.name AS name, person.gender AS gender, person.alias AS alias\",\n",
    "                                    person_name=person_name)\n",
    "            record = result.single()\n",
    "            if record:\n",
    "                return {\n",
    "                    \"name\": record[\"name\"],\n",
    "                    \"gender\": record[\"gender\"],\n",
    "                    \"alias\": record[\"alias\"]\n",
    "                }\n",
    "            else:\n",
    "                return None\n",
    "\n",
    "\n",
    "uri      = \"neo4j+s://5bb87071.databases.neo4j.io\"\n",
    "username = \"neo4j\"\n",
    "password = \"PPcEmlN5ognHV-YG0J9Puj0FzCwdiXsElSbaHugN9xA\"\n",
    "\n",
    "manager = Neo4jManager(uri, username, password)\n",
    "manager.clear_database()\n",
    "manager.import_characters('./Persons.csv')\n",
    "manager.import_relationships('./Relationship.csv')\n",
    "status = manager.query_status()\n",
    "print(status)\n",
    "\n",
    "person_name = '甄嬛'\n",
    "relationship_type = '父亲'\n",
    "related_names = manager.query_relationship(person_name, relationship_type)\n",
    "if related_names:\n",
    "    print(f\"{person_name}的{relationship_type}是: {', '.join(related_names)}\")\n",
    "else:\n",
    "    print(f\"{person_name}没有{relationship_type}\")\n",
    "\n",
    "aliases = manager.query_alias(person_name)\n",
    "if aliases:\n",
    "    print(f\"{person_name}的别名: {aliases}\")\n",
    "else:\n",
    "    print(f\"{person_name}没有别名\")\n",
    "\n",
    "person_info = manager.query_person_info(person_name)\n",
    "if person_info:\n",
    "    print(\"人物信息:\")\n",
    "    print(f\"姓名: {person_info['name']}\")\n",
    "    print(f\"性别: {person_info['gender']}\")\n",
    "    print(f\"别名: {person_info['alias']}\")\n",
    "else:\n",
    "    print(f\"{person_name}不存在\")\n",
    "\n",
    "manager.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## langchain\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "        Node properties are the following:\n",
      "        [{'properties': [{'property': 'gender', 'type': 'STRING'}, {'property': 'name', 'type': 'STRING'}, {'property': 'alias', 'type': 'STRING'}], 'labels': 'Character'}]\n",
      "        Relationship properties are the following:\n",
      "        []\n",
      "        The relationships are the following:\n",
      "        ['(:Character)-[:爱人]->(:Character)', '(:Character)-[:朋友]->(:Character)', '(:Character)-[:女儿]->(:Character)', '(:Character)-[:父亲]->(:Character)', '(:Character)-[:妻子]->(:Character)', '(:Character)-[:妃子]->(:Character)', '(:Character)-[:丈夫]->(:Character)', '(:Character)-[:敌人]->(:Character)']\n",
      "        \n"
     ]
    }
   ],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.chains import GraphCypherQAChain\n",
    "from langchain.graphs import Neo4jGraph\n",
    "\n",
    "uri      = \"neo4j+s://5bb87071.databases.neo4j.io\"\n",
    "username = \"neo4j\"\n",
    "password = \"PPcEmlN5ognHV-YG0J9Puj0FzCwdiXsElSbaHugN9xA\"\n",
    "\n",
    "graph = Neo4jGraph(\n",
    "    url=uri, username=username, password=password\n",
    ")\n",
    "\n",
    "# graph.refresh_schema()\n",
    "print(graph.get_schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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\n",
      "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\n"
     ]
    }
   ],
   "source": [
    "chain = GraphCypherQAChain.from_llm(\n",
    "    ChatOpenAI(temperature=0), graph=graph, verbose=True\n",
    ")\n",
    "\n",
    "# openai.api_key = 'sk-5vOFyAXiUcLGkRKV60D2D7E6Ee6e408f8d366f4c1013C004'\n",
    "# # 更换OpenAI接口的host\n",
    "# openai.api_base = \"https://api.akm.pw/v1\" #在这里设置即可,需要特别注意这里的/v1是必须的，否则报错。前面的地址注意替换即可。\n",
    "\n",
    "import os \n",
    "from dotenv import load_dotenv, find_dotenv, dotenv_values \n",
    "# os.environ[\"OPENAI_API_KEY\"] = \"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\"\n",
    "os.environ['HTTPS_PROXY']    = 'http://127.0.0.1:7890'\n",
    "os.environ[\"HTTP_PROXY\"]     = 'http://127.0.0.1:7890'\n",
    "\n",
    "# 读取本地的环境变量 \n",
    "env_vars = dotenv_values('.env')\n",
    "# 获取环境变量 OPENAI_API_KEY\n",
    "openai_api_key = env_vars['OPENAI_API_KEY']\n",
    "print(openai_api_key)\n",
    "\n",
    "# 获取环境变量的值\n",
    "api_key = os.environ.get('OPENAI_API_KEY')\n",
    "\n",
    "# 打印环境变量的值\n",
    "print(api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mMATCH (c1:Character {name: '甄嬛'})-[:女儿]->(c2:Character)\n",
      "RETURN c2.name\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'c2.name': '胧月公主'}]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'甄嬛的女儿是胧月公主。'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(\"甄嬛的女儿是谁？\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mMATCH (queen:Character {name: '皇后'})-[:敌人]->(enemy:Character)\n",
      "RETURN enemy.name\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'enemy.name': '甄嬛'}]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'皇后的敌人是甄嬛。'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(\"皇后的敌人是谁？\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mMATCH (zhenhuan:Character {name: '甄嬛'})-[:女儿]->(child:Character)\n",
      "RETURN child.name\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'child.name': '胧月公主'}]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'甄嬛的孩子是胧月公主。'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(\"甄嬛的孩子是谁？\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch_gpu",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
